// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/exec/aggregate.h>
#include <pollux/exec/simple_aggregate_adapter.h>
#include <pollux/expression/function_signature.h>
#include <pollux/expression/vector_writers.h>
#include <pollux/functions/lib/aggregates/value_list.h>

using namespace kumo::pollux::exec;

namespace kumo::pollux::aggregate {

namespace {
class ArrayAggAggregate {
 public:
  // Type(s) of input vector(s) wrapped in Row.
  using InputType = Row<Generic<T1>>;

  // Type of intermediate result vector.
  using IntermediateType = Array<Generic<T1>>;

  // Type of output vector.
  using OutputType = Array<Generic<T1>>;

  static constexpr bool default_null_behavior_ = false;

  static bool toIntermediate(
      exec::out_type<Array<Generic<T1>>>& out,
      exec::optional_arg_type<Generic<T1>> in) {
    if (in.has_value()) {
      out.add_item().copy_from(in.value());
    } else {
      out.add_null();
    }
    return true;
  }

  struct AccumulatorType {
    ValueList elements_;

    AccumulatorType() = delete;

    // Constructor used in initializeNewGroups().
    explicit AccumulatorType(
        HashStringAllocator* /*allocator*/,
        ArrayAggAggregate* /*fn*/)
        : elements_{} {}

    static constexpr bool is_fixed_size_ = false;

    // addInput expects one parameter of exec::optional_arg_type<T> for each
    // child-type T wrapped in InputType.
    bool addInput(
        HashStringAllocator* allocator,
        exec::optional_arg_type<Generic<T1>> data) {
      elements_.appendValue(data, allocator);
      return true;
    }

    // combine expects one parameter of
    // exec::optional_arg_type<IntermediateType>.
    bool combine(
        HashStringAllocator* allocator,
        exec::optional_arg_type<Array<Generic<T1>>> other) {
      if (!other.has_value()) {
        return false;
      }
      for (auto element : other.value()) {
        elements_.appendValue(element, allocator);
      }
      return true;
    }

    bool writeFinalResult(
        bool nonNullGroup,
        exec::out_type<Array<Generic<T1>>>& out) {
      if (!nonNullGroup) {
        return false;
      }
      copyValueListToArrayWriter(out, elements_);
      return true;
    }

    bool writeIntermediateResult(
        bool nonNullGroup,
        exec::out_type<Array<Generic<T1>>>& out) {
      // If the group's accumulator is null, the corresponding intermediate
      // result is null too.
      if (!nonNullGroup) {
        return false;
      }
      copyValueListToArrayWriter(out, elements_);
      return true;
    }

    void destroy(HashStringAllocator* allocator) {
      elements_.free(allocator);
    }
  };
};

} // namespace

exec::AggregateRegistrationResult registerSimpleArrayAggAggregate(
    const std::string& name) {
  std::vector<std::shared_ptr<exec::AggregateFunctionSignature>> signatures{
      exec::AggregateFunctionSignatureBuilder()
          .typeVariable("E")
          .returnType("array(E)")
          .intermediateType("array(E)")
          .argumentType("E")
          .build()};

  return exec::registerAggregateFunction(
      name,
      std::move(signatures),
      [name](
          core::AggregationNode::Step step,
          const std::vector<TypePtr>& argTypes,
          const TypePtr& resultType,
          const core::QueryConfig& /*config*/)
          -> std::unique_ptr<exec::Aggregate> {
        POLLUX_CHECK_EQ(
            argTypes.size(), 1, "{} takes at most one argument", name);
        return std::make_unique<SimpleAggregateAdapter<ArrayAggAggregate>>(
            step, argTypes, resultType);
      },
      true /*registerCompanionFunctions*/,
      true /*overwrite*/);
}

} // namespace kumo::pollux::aggregate
